#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Pspy Integration Module
Pspy进程监控集成模块
"""

import os
import re
import json
import time
import signal
import threading
import subprocess
from typing import Dict, List, Any, Optional, Callable
from pathlib import Path
from .tool_manager import ExternalToolManager

class PspyIntegration:
    """Pspy集成类"""
    
    def __init__(self, tool_manager: ExternalToolManager = None):
        self.tool_manager = tool_manager or ExternalToolManager()
        self.tool_name = "pspy"
        self.monitoring_process = None
        self.monitoring_thread = None
        self.is_monitoring = False
        self.captured_processes = []
        self.process_callbacks = []
        self.monitoring_start_time = None
        self.monitoring_duration = 0
    
    def is_available(self) -> bool:
        """检查pspy是否可用"""
        status = self.tool_manager.get_tool_status(self.tool_name)
        return status == "ready"
    
    def ensure_available(self) -> bool:
        """确保pspy可用"""
        if self.is_available():
            return True
        
        print("📥 pspy不可用，正在下载...")
        return self.tool_manager.download_tool(self.tool_name)
    
    def start_monitoring(self, duration: int = 300, interval: int = 1000, 
                        include_fs_events: bool = True, 
                        include_process_events: bool = True) -> bool:
        """开始进程监控"""
        if not self.ensure_available():
            print("❌ pspy不可用")
            return False
        
        if self.is_monitoring:
            print("⚠️ 监控已在运行中")
            return False
        
        # 构造参数
        args = []
        
        # 设置间隔
        args.extend(["-i", str(interval)])
        
        # 文件系统事件
        if not include_fs_events:
            args.append("-f")
        
        # 进程事件
        if not include_process_events:
            args.append("-p")
        
        try:
            print(f"🔍 开始pspy监控 (持续时间: {duration}秒)")
            self.monitoring_start_time = time.time()
            self.monitoring_duration = duration
            self.captured_processes = []
            self.is_monitoring = True
            
            # 启动监控线程
            self.monitoring_thread = threading.Thread(
                target=self._monitoring_worker,
                args=(args, duration),
                daemon=True
            )
            self.monitoring_thread.start()
            
            return True
            
        except Exception as e:
            print(f"❌ 启动监控失败: {e}")
            self.is_monitoring = False
            return False
    
    def _monitoring_worker(self, args: List[str], duration: int):
        """监控工作线程"""
        try:
            # 获取pspy路径
            pspy_path = self.tool_manager.get_tool_path(self.tool_name)
            if not pspy_path:
                print("❌ 无法获取pspy路径")
                return
            
            # 启动pspy进程
            cmd = [pspy_path] + args
            
            self.monitoring_process = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                bufsize=1,
                universal_newlines=True
            )
            
            print(f"✅ pspy监控已启动 (PID: {self.monitoring_process.pid})")
            
            # 读取输出
            start_time = time.time()
            
            while self.is_monitoring and time.time() - start_time < duration:
                try:
                    # 检查进程是否还在运行
                    if self.monitoring_process.poll() is not None:
                        break
                    
                    # 读取一行输出
                    line = self.monitoring_process.stdout.readline()
                    if line:
                        self._process_output_line(line.strip())
                    
                    time.sleep(0.1)  # 短暂休眠
                    
                except Exception as e:
                    print(f"⚠️ 读取输出错误: {e}")
                    break
            
            # 停止监控
            self._stop_monitoring_process()
            
        except Exception as e:
            print(f"❌ 监控工作线程错误: {e}")
        finally:
            self.is_monitoring = False
            print("🛑 pspy监控已停止")
    
    def _process_output_line(self, line: str):
        """处理输出行"""
        try:
            # 解析pspy输出格式
            process_info = self._parse_pspy_line(line)
            
            if process_info:
                # 添加时间戳
                process_info["timestamp"] = time.time()
                process_info["relative_time"] = time.time() - self.monitoring_start_time
                
                # 保存到捕获列表
                self.captured_processes.append(process_info)
                
                # 调用回调函数
                for callback in self.process_callbacks:
                    try:
                        callback(process_info)
                    except Exception as e:
                        print(f"⚠️ 回调函数错误: {e}")
                
                # 实时显示重要进程
                if self._is_interesting_process(process_info):
                    self._display_interesting_process(process_info)
                    
        except Exception as e:
            print(f"⚠️ 处理输出行错误: {e}")
    
    def _parse_pspy_line(self, line: str) -> Optional[Dict[str, Any]]:
        """解析pspy输出行"""
        try:
            # pspy输出格式示例:
            # 2023-10-12 14:30:15.123 UID=0    PID=1234   | /bin/bash -c "ls -la"
            
            # 匹配时间戳
            timestamp_pattern = r'(\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}\.\d{3})'
            timestamp_match = re.search(timestamp_pattern, line)
            
            if not timestamp_match:
                return None
            
            timestamp_str = timestamp_match.group(1)
            
            # 匹配UID和PID
            uid_pattern = r'UID=(\d+)'
            pid_pattern = r'PID=(\d+)'
            
            uid_match = re.search(uid_pattern, line)
            pid_match = re.search(pid_pattern, line)
            
            if not (uid_match and pid_match):
                return None
            
            uid = int(uid_match.group(1))
            pid = int(pid_match.group(1))
            
            # 提取命令
            command_start = line.find('|')
            if command_start == -1:
                return None
            
            command = line[command_start + 1:].strip()
            
            return {
                "timestamp_str": timestamp_str,
                "uid": uid,
                "pid": pid,
                "command": command,
                "full_line": line,
                "is_root": uid == 0,
                "command_parts": command.split() if command else []
            }
            
        except Exception as e:
            print(f"⚠️ 解析行错误: {e}")
            return None
    
    def _is_interesting_process(self, process_info: Dict[str, Any]) -> bool:
        """判断是否为有趣的进程"""
        command = process_info.get("command", "").lower()
        uid = process_info.get("uid", -1)
        
        # 检查有趣的命令
        interesting_commands = [
            "sudo", "su", "passwd", "ssh", "scp", "rsync",
            "wget", "curl", "nc", "netcat", "nmap",
            "python", "perl", "ruby", "php", "bash", "sh",
            "crontab", "systemctl", "service",
            "mount", "umount", "chmod", "chown",
            "iptables", "ufw", "firewall"
        ]
        
        # 检查有趣的路径
        interesting_paths = [
            "/tmp/", "/var/tmp/", "/dev/shm/",
            "/etc/", "/root/", "/home/",
            ".ssh/", ".bash_history", ".mysql_history"
        ]
        
        # 检查命令
        for cmd in interesting_commands:
            if cmd in command:
                return True
        
        # 检查路径
        for path in interesting_paths:
            if path in command:
                return True
        
        # 检查root进程
        if uid == 0:
            return True
        
        return False
    
    def _display_interesting_process(self, process_info: Dict[str, Any]):
        """显示有趣的进程"""
        uid = process_info.get("uid", -1)
        pid = process_info.get("pid", -1)
        command = process_info.get("command", "")
        relative_time = process_info.get("relative_time", 0)
        
        user_indicator = "🔴" if uid == 0 else "🟡"
        
        print(f"{user_indicator} [{relative_time:.1f}s] UID={uid} PID={pid} | {command[:80]}{'...' if len(command) > 80 else ''}")
    
    def stop_monitoring(self) -> bool:
        """停止监控"""
        if not self.is_monitoring:
            print("⚠️ 监控未在运行")
            return False
        
        print("🛑 正在停止pspy监控...")
        self.is_monitoring = False
        
        # 停止监控进程
        self._stop_monitoring_process()
        
        # 等待线程结束
        if self.monitoring_thread and self.monitoring_thread.is_alive():
            self.monitoring_thread.join(timeout=5)
        
        print("✅ pspy监控已停止")
        return True
    
    def _stop_monitoring_process(self):
        """停止监控进程"""
        if self.monitoring_process:
            try:
                # 尝试优雅地终止进程
                self.monitoring_process.terminate()
                
                # 等待进程结束
                try:
                    self.monitoring_process.wait(timeout=5)
                except subprocess.TimeoutExpired:
                    # 强制杀死进程
                    self.monitoring_process.kill()
                    self.monitoring_process.wait()
                
            except Exception as e:
                print(f"⚠️ 停止监控进程错误: {e}")
            finally:
                self.monitoring_process = None
    
    def get_monitoring_status(self) -> Dict[str, Any]:
        """获取监控状态"""
        return {
            "is_monitoring": self.is_monitoring,
            "start_time": self.monitoring_start_time,
            "duration": self.monitoring_duration,
            "captured_processes_count": len(self.captured_processes),
            "monitoring_pid": self.monitoring_process.pid if self.monitoring_process else None
        }
    
    def get_captured_processes(self, filter_interesting: bool = False) -> List[Dict[str, Any]]:
        """获取捕获的进程"""
        if filter_interesting:
            return [p for p in self.captured_processes if self._is_interesting_process(p)]
        else:
            return self.captured_processes.copy()
    
    def search_processes(self, keyword: str, case_sensitive: bool = False) -> List[Dict[str, Any]]:
        """搜索进程"""
        results = []
        
        for process in self.captured_processes:
            command = process.get("command", "")
            
            if case_sensitive:
                match = keyword in command
            else:
                match = keyword.lower() in command.lower()
            
            if match:
                results.append(process)
        
        return results
    
    def get_process_statistics(self) -> Dict[str, Any]:
        """获取进程统计"""
        if not self.captured_processes:
            return {"error": "没有捕获的进程数据"}
        
        total_processes = len(self.captured_processes)
        root_processes = len([p for p in self.captured_processes if p.get("uid") == 0])
        interesting_processes = len([p for p in self.captured_processes if self._is_interesting_process(p)])
        
        # 统计最常见的命令
        command_counts = {}
        for process in self.captured_processes:
            cmd_parts = process.get("command_parts", [])
            if cmd_parts:
                base_cmd = cmd_parts[0]
                command_counts[base_cmd] = command_counts.get(base_cmd, 0) + 1
        
        # 获取前10个最常见的命令
        top_commands = sorted(command_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        
        # 统计UID分布
        uid_counts = {}
        for process in self.captured_processes:
            uid = process.get("uid", -1)
            uid_counts[uid] = uid_counts.get(uid, 0) + 1
        
        return {
            "total_processes": total_processes,
            "root_processes": root_processes,
            "interesting_processes": interesting_processes,
            "root_percentage": (root_processes / total_processes * 100) if total_processes > 0 else 0,
            "interesting_percentage": (interesting_processes / total_processes * 100) if total_processes > 0 else 0,
            "top_commands": top_commands,
            "uid_distribution": uid_counts,
            "monitoring_duration": time.time() - self.monitoring_start_time if self.monitoring_start_time else 0
        }
    
    def export_results(self, file_path: str, format: str = "json", 
                      filter_interesting: bool = False):
        """导出结果"""
        if not self.captured_processes:
            print("❌ 没有可导出的进程数据")
            return False
        
        try:
            processes = self.get_captured_processes(filter_interesting)
            
            if format.lower() == "json":
                export_data = {
                    "monitoring_info": {
                        "start_time": self.monitoring_start_time,
                        "duration": self.monitoring_duration,
                        "total_captured": len(self.captured_processes),
                        "exported_count": len(processes)
                    },
                    "statistics": self.get_process_statistics(),
                    "processes": processes
                }
                
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(export_data, f, indent=2, ensure_ascii=False)
            
            elif format.lower() == "txt":
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("Pspy进程监控结果\n")
                    f.write("=" * 50 + "\n\n")
                    
                    stats = self.get_process_statistics()
                    f.write(f"监控时长: {stats.get('monitoring_duration', 0):.1f}秒\n")
                    f.write(f"总进程数: {stats.get('total_processes', 0)}\n")
                    f.write(f"Root进程: {stats.get('root_processes', 0)} ({stats.get('root_percentage', 0):.1f}%)\n")
                    f.write(f"有趣进程: {stats.get('interesting_processes', 0)} ({stats.get('interesting_percentage', 0):.1f}%)\n\n")
                    
                    f.write("捕获的进程:\n")
                    f.write("-" * 30 + "\n")
                    
                    for process in processes:
                        f.write(f"[{process.get('relative_time', 0):.1f}s] ")
                        f.write(f"UID={process.get('uid', -1)} ")
                        f.write(f"PID={process.get('pid', -1)} | ")
                        f.write(f"{process.get('command', '')}\n")
            
            print(f"📄 结果已导出到: {file_path}")
            return True
            
        except Exception as e:
            print(f"❌ 导出失败: {e}")
            return False
    
    def add_process_callback(self, callback: Callable[[Dict[str, Any]], None]):
        """添加进程回调函数"""
        self.process_callbacks.append(callback)
    
    def remove_process_callback(self, callback: Callable[[Dict[str, Any]], None]):
        """移除进程回调函数"""
        if callback in self.process_callbacks:
            self.process_callbacks.remove(callback)
    
    def clear_captured_processes(self):
        """清空捕获的进程"""
        self.captured_processes = []
        print("🗑️ 已清空捕获的进程数据")

# 使用示例
def example_usage():
    """使用示例"""
    # 创建pspy集成实例
    pspy = PspyIntegration()
    
    # 检查可用性
    if pspy.is_available():
        print("✅ pspy已就绪")
    else:
        print("📥 正在准备pspy...")
        if not pspy.ensure_available():
            print("❌ pspy准备失败")
            return
    
    # 添加进程回调
    def process_callback(process_info):
        if process_info.get("uid") == 0:  # 只显示root进程
            print(f"🔴 Root进程: {process_info.get('command', '')[:50]}")
    
    pspy.add_process_callback(process_callback)
    
    # 开始监控
    print("🔍 开始进程监控...")
    if pspy.start_monitoring(duration=60, interval=500):  # 监控60秒，间隔500ms
        
        # 等待一段时间
        time.sleep(30)
        
        # 获取当前状态
        status = pspy.get_monitoring_status()
        print(f"📊 监控状态: 已捕获 {status['captured_processes_count']} 个进程")
        
        # 搜索特定进程
        sudo_processes = pspy.search_processes("sudo")
        if sudo_processes:
            print(f"🔍 找到 {len(sudo_processes)} 个sudo相关进程")
        
        # 等待监控完成
        time.sleep(30)
        
        # 获取统计信息
        stats = pspy.get_process_statistics()
        print(f"📊 监控统计:")
        print(f"  - 总进程: {stats['total_processes']}")
        print(f"  - Root进程: {stats['root_processes']} ({stats['root_percentage']:.1f}%)")
        print(f"  - 有趣进程: {stats['interesting_processes']} ({stats['interesting_percentage']:.1f}%)")
        
        # 导出结果
        pspy.export_results("pspy_results.json", "json", filter_interesting=True)
        
    else:
        print("❌ 监控启动失败")

if __name__ == "__main__":
    example_usage()